{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# extern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['PATH'] += ':/usr/local/cuda/bin' # 保证 nvcc 可以被找到\n",
    "import tvm\n",
    "from tvm import te\n",
    "import numpy as np\n",
    "import tvm.testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "验证 TVM 在不同目标设备上的向量化代码生成能力。测试分为 CPU 和 GPU 两个版本：CPU 版本使用循环展开策略处理向量化计算，GPU 版本通过线程块和线程索引实现并行。核心逻辑通过手动构建 TIR 中间表示，验证生成代码在 LLVM/OpenCL/CUDA 后端的正确性。测试使用 `te.extern` 创建外部计算节点，并检查输出结果是否符合预期。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CPU 版本：使用 SIMD 向量化策略（`float32x2`），每次迭代处理 2 个元素，实现 2 倍循环展开"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extern_generator(ins, outs):\n",
    "    \"\"\"Manually write the IR for the extern function, add pipeline\"\"\"\n",
    "    ib = tvm.tir.ir_builder.create()\n",
    "    with ib.for_range(0, (n + 1) // 2) as i:\n",
    "        ib.emit(\n",
    "            outs[0].vstore(\n",
    "                i * 2, ins[0].vload(i * 2, \"float32x2\") + tvm.tir.const(1, \"float32x2\")\n",
    "            )\n",
    "        )\n",
    "    return ib.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GPU 版本：通过 `blockIdx.x` 和 `threadIdx.x` 实现两级并行，适配 GPU 的 SIMT 架构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extern_generator_gpu(ins, outs):\n",
    "    \"\"\"Manually write the IR for the extern function, add pipeline\"\"\"\n",
    "    ib = tvm.tir.ir_builder.create()\n",
    "    bx = te.thread_axis(\"blockIdx.x\")\n",
    "    tx = te.thread_axis(\"threadIdx.x\")\n",
    "    ib.scope_attr(bx, \"thread_extent\", (nn + max_threads - 1) // max_threads)\n",
    "    ib.scope_attr(tx, \"thread_extent\", max_threads)\n",
    "    idx = bx.var * max_threads + tx.var\n",
    "    with ib.if_scope(ib.likely(idx < n)):\n",
    "        ib.emit(\n",
    "            outs[0].vstore(\n",
    "                idx * 2, ins[0].vload(idx * 2, \"float32x2\") + tvm.tir.const(1, \"float32x2\")\n",
    "            )\n",
    "        )\n",
    "    return ib.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- `te.extern` 创建外部计算节点，分离计算定义与实现\n",
    "- `vload/vstore` 实现显式向量化内存访问\n",
    "- 内存对齐：向量化访问要求 64 位对齐（`float32x2`对应 `2*4B=8B`）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn = 64\n",
    "max_threads = 4\n",
    "n = tvm.runtime.convert(nn)\n",
    "A = te.placeholder((n,), name=\"A\")\n",
    "\n",
    "C_cpu = te.extern(A.shape, [A], extern_generator, name=\"C\")\n",
    "C_gpu = te.extern(A.shape, [A], extern_generator_gpu, name=\"C\")\n",
    "\n",
    "# Create IRModules directly\n",
    "mod_cpu = tvm.IRModule.from_expr(te.create_prim_func([A, C_cpu]))\n",
    "mod_gpu = tvm.IRModule.from_expr(te.create_prim_func([A, C_gpu]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(var_A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        A <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_A, (<span style=\"color: #008000\">64</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        C <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_C, (<span style=\"color: #008000\">64</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes()\n",
       "            <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">32</span>):\n",
       "                C[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span>:i <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span>:i <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Broadcast(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1.0</span>), <span style=\"color: #008000\">2</span>)\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod_cpu.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(var_A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        A <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_A, (<span style=\"color: #008000\">64</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        C <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_C, (<span style=\"color: #008000\">64</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes()\n",
       "            blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>launch_thread(<span style=\"color: #BA2121\">&quot;blockIdx.x&quot;</span>, <span style=\"color: #008000\">16</span>)\n",
       "            threadIdx_x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>launch_thread(<span style=\"color: #BA2121\">&quot;threadIdx.x&quot;</span>, <span style=\"color: #008000\">4</span>)\n",
       "            <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> <span style=\"color: #008000\">64</span>):\n",
       "                C[(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span>:(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span>:(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Broadcast(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1.0</span>), <span style=\"color: #008000\">2</span>)\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod_gpu.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "跨设备统一验证："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_target(target):\n",
    "    if not tvm.testing.device_enabled(target):\n",
    "        return\n",
    "    mod = mod_gpu if target in [\"opencl\", \"cuda\"] else mod_cpu\n",
    "    C = C_gpu if target in [\"opencl\", \"cuda\"] else C_cpu\n",
    "    # build and invoke the kernel.\n",
    "    f = tvm.compile(mod, target=target)\n",
    "    dev = tvm.device(target, 0)\n",
    "    # launch the kernel.\n",
    "    n = nn\n",
    "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
    "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
    "    f(a, c)\n",
    "    tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_target(\"llvm\")\n",
    "check_target(\"opencl\")\n",
    "check_target(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 打包 buffer "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extern_generator(ins, outs):\n",
    "    \"\"\"Manually write the IR for the extern function, add pipeline.\"\"\"\n",
    "    return tvm.tir.call_packed(\"my_extern_array_func1\", ins[0], outs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(var_A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, var_C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        A <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_A, (<span style=\"color: #008000\">1024</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        C <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_C, (<span style=\"color: #008000\">1024</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_packed(<span style=\"color: #BA2121\">&quot;my_extern_array_func1&quot;</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_array(A<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_shape(<span style=\"color: #008000\">1024</span>), <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>), A<span style=\"color: #AA22FF; font-weight: bold\">.</span>elem_offset), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_array(C<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_shape(<span style=\"color: #008000\">1024</span>), <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>), C<span style=\"color: #AA22FF; font-weight: bold\">.</span>elem_offset))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nn = 1024\n",
    "n = tvm.runtime.convert(nn)\n",
    "A = te.placeholder((n,), name=\"A\")\n",
    "C = te.extern(A.shape, [A], extern_generator, name=\"C\")\n",
    "\n",
    "# Create IRModule directly\n",
    "mod = tvm.IRModule.from_expr(te.create_prim_func([A, C]))\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.register_func\n",
    "def my_extern_array_func1(aa, bb):\n",
    "    aa.copyto(bb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_target(target):\n",
    "    if not tvm.testing.device_enabled(target):\n",
    "        return\n",
    "    # build and invoke the kernel.\n",
    "    f = tvm.compile(mod, target=target)\n",
    "    dev = tvm.cpu(0)\n",
    "    # launch the kernel.\n",
    "    n = nn\n",
    "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
    "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
    "\n",
    "    f(a, c)\n",
    "    tvm.testing.assert_allclose(c.numpy(), a.numpy())\n",
    "\n",
    "check_target(\"llvm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 打包缓冲区中间表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extern_generator(ins, outs):\n",
    "    \"\"\"Manually write the IR for the extern function, add pipeline.\"\"\"\n",
    "    return tvm.tir.call_packed(\"my_extern_array_func2\", *ins, outs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "\n",
       "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), var_C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        C <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(var_C, (<span style=\"color: #008000\">1024</span>,), offset_factor<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>)\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        B <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer((<span style=\"color: #008000\">1024</span>,))\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">1024</span>):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;B&quot;</span>):\n",
       "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, i)\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(B[v_i])\n",
       "                B[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1.0</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes()\n",
       "            elem_offset <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
       "            T<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_packed(<span style=\"color: #BA2121\">&quot;my_extern_array_func2&quot;</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_array(B<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_shape(<span style=\"color: #008000\">1024</span>), <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>), elem_offset), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_array(C<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>tvm_stack_make_shape(<span style=\"color: #008000\">1024</span>), <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">1</span>, T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0.0</span>), C<span style=\"color: #AA22FF; font-weight: bold\">.</span>elem_offset))\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nn = 1024\n",
    "n = tvm.runtime.convert(nn)\n",
    "A = te.placeholder((n,), name=\"A\")\n",
    "B = te.compute((n,), lambda i: A[i] + 1, name=\"B\")\n",
    "C = te.extern(B.shape, [B], extern_generator, name=\"C\")\n",
    "# D = te.compute((n,), lambda i: C[i] + 1, name=\"D\")\n",
    "mod = tvm.IRModule.from_expr(te.create_prim_func([A, C]))\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_target(target):\n",
    "    if not tvm.testing.device_enabled(target):\n",
    "        return\n",
    "    # build and invoke the kernel.\n",
    "    f = tvm.compile(mod, target=target)\n",
    "    dev = tvm.cpu(0)\n",
    "    # launch the kernel.\n",
    "    n = nn\n",
    "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
    "    b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev)\n",
    "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
    "\n",
    "    @tvm.register_func\n",
    "    def my_extern_array_func2(aa, cc):\n",
    "        assert aa.shape == a.shape\n",
    "        tvm.testing.assert_allclose(aa.numpy(), a.numpy()+1)\n",
    "        aa.copyto(cc)\n",
    "\n",
    "    f(a, c)\n",
    "    tvm.testing.assert_allclose(c.numpy(), a.numpy() + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "InternalError",
     "evalue": "Check failed: undefined.size() == 0 (1 vs. 0) : In PrimFunc main variables [elem_offset] are used, but are not passed in as API arguments",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInternalError\u001b[0m                             Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[17], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mcheck_target\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mllvm\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[16], line 5\u001b[0m, in \u001b[0;36mcheck_target\u001b[0;34m(target)\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;66;03m# build and invoke the kernel.\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m f \u001b[38;5;241m=\u001b[39m \u001b[43mtvm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmod\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m dev \u001b[38;5;241m=\u001b[39m tvm\u001b[38;5;241m.\u001b[39mcpu(\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m# launch the kernel.\u001b[39;00m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/driver/build_module.py:110\u001b[0m, in \u001b[0;36mcompile\u001b[0;34m(mod, target, relax_pipeline, tir_pipeline)\u001b[0m\n\u001b[1;32m    103\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _contains_relax(mod):\n\u001b[1;32m    104\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m tvm\u001b[38;5;241m.\u001b[39mrelax\u001b[38;5;241m.\u001b[39mbuild(\n\u001b[1;32m    105\u001b[0m         mod,\n\u001b[1;32m    106\u001b[0m         target,\n\u001b[1;32m    107\u001b[0m         relax_pipeline\u001b[38;5;241m=\u001b[39mrelax_pipeline,\n\u001b[1;32m    108\u001b[0m         tir_pipeline\u001b[38;5;241m=\u001b[39mtir_pipeline,\n\u001b[1;32m    109\u001b[0m     )\n\u001b[0;32m--> 110\u001b[0m lib \u001b[38;5;241m=\u001b[39m \u001b[43mtvm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtir\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbuild\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmod\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpipeline\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtir_pipeline\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m Executable(lib)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/tir/build.py:173\u001b[0m, in \u001b[0;36mbuild\u001b[0;34m(mod, target, pipeline)\u001b[0m\n\u001b[1;32m    170\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    171\u001b[0m     \u001b[38;5;66;03m# default pipeline depends on the target\u001b[39;00m\n\u001b[1;32m    172\u001b[0m     pipeline \u001b[38;5;241m=\u001b[39m tvm\u001b[38;5;241m.\u001b[39mtir\u001b[38;5;241m.\u001b[39mget_default_tir_pipeline(target)\n\u001b[0;32m--> 173\u001b[0m mod \u001b[38;5;241m=\u001b[39m \u001b[43mpipeline\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmod\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    175\u001b[0m \u001b[38;5;66;03m# Step 5: Get host and device modules\u001b[39;00m\n\u001b[1;32m    176\u001b[0m host_mod, device_mod_dict \u001b[38;5;241m=\u001b[39m split_host_device_mods(mod)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/ir/transform.py:238\u001b[0m, in \u001b[0;36mPass.__call__\u001b[0;34m(self, mod)\u001b[0m\n\u001b[1;32m    224\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, mod):\n\u001b[1;32m    225\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Execute the pass. Note that for sequential pass, the dependency among\u001b[39;00m\n\u001b[1;32m    226\u001b[0m \u001b[38;5;124;03m    different passes will be resolved in the backend.\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    236\u001b[0m \u001b[38;5;124;03m        The updated module after applying this pass.\u001b[39;00m\n\u001b[1;32m    237\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 238\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ffi_transform_api\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mRunPass\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmod\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/ffi/cython/function.pxi:228\u001b[0m, in \u001b[0;36mtvm.ffi.core.Function.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:576\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    574\u001b[0m \n\u001b[1;32m    575\u001b[0m TVM_REGISTER_GLOBAL(\"transform.RunPass\")\n\u001b[0;32m--> 576\u001b[0m     .set_body_typed([](Pass pass, ffi::RValueRef<IRModule> mod) { return pass(*std::move(mod)); });\n\u001b[1;32m    577\u001b[0m \n\u001b[1;32m    578\u001b[0m TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:297\u001b[0m, in \u001b[0;36mtvm::transform::Pass::operator()(tvm::IRModule) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    295\u001b[0m \n\u001b[1;32m    296\u001b[0m IRModule Pass::operator()(IRModule mod) const {\n\u001b[0;32m--> 297\u001b[0m   return this->operator()(std::move(mod), PassContext::Current());\n\u001b[1;32m    298\u001b[0m }\n\u001b[1;32m    299\u001b[0m \n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:313\u001b[0m, in \u001b[0;36mtvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    311\u001b[0m   ret = Pass::AssertImmutableModule(mod, node, pass_ctx);\n\u001b[1;32m    312\u001b[0m } else {\n\u001b[0;32m--> 313\u001b[0m   ret = node->operator()(std::move(mod), pass_ctx);\n\u001b[1;32m    314\u001b[0m }\n\u001b[1;32m    315\u001b[0m pass_ctx.InstrumentAfterPass(ret, pass_info);\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:419\u001b[0m, in \u001b[0;36mtvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    417\u001b[0m VLOG(0) << \"Executing module pass with opt level: \" << pass_info->opt_level;\n\u001b[1;32m    418\u001b[0m \n\u001b[0;32m--> 419\u001b[0m mod = pass_func(std::move(mod), pass_ctx);\n\u001b[1;32m    420\u001b[0m \n\u001b[1;32m    421\u001b[0m ICHECK(mod.defined()) << \"The return value of a module pass must be set.\";\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:570\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    568\u001b[0m  PassInfo pass_info) {\n\u001b[1;32m    569\u001b[0m auto wrapped_pass_func = [pass_func](IRModule mod, PassContext ctx) {\n\u001b[0;32m--> 570\u001b[0m   return pass_func(ffi::RValueRef<IRModule>(std::move(mod)), ctx);\n\u001b[1;32m    571\u001b[0m };\n\u001b[1;32m    572\u001b[0m return ModulePass(wrapped_pass_func, pass_info);\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/ffi/cython/function.pxi:281\u001b[0m, in \u001b[0;36mtvm.ffi.core.tvm_ffi_callback\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/tir/pipeline.py:122\u001b[0m, in \u001b[0;36m_pipeline\u001b[0;34m()\u001b[0m\n\u001b[1;32m    109\u001b[0m     passes\u001b[38;5;241m.\u001b[39mappend(tir\u001b[38;5;241m.\u001b[39mtransform\u001b[38;5;241m.\u001b[39mInjectPTXLDG32())\n\u001b[1;32m    110\u001b[0m passes\u001b[38;5;241m.\u001b[39mextend(\n\u001b[1;32m    111\u001b[0m     [\n\u001b[1;32m    112\u001b[0m         tir\u001b[38;5;241m.\u001b[39mtransform\u001b[38;5;241m.\u001b[39mAnnotateDeviceRegions(),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    120\u001b[0m     ]\n\u001b[1;32m    121\u001b[0m )\n\u001b[0;32m--> 122\u001b[0m mod \u001b[38;5;241m=\u001b[39m tvm\u001b[38;5;241m.\u001b[39mir\u001b[38;5;241m.\u001b[39mtransform\u001b[38;5;241m.\u001b[39mSequential(passes)(mod)\n\u001b[1;32m    123\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m mod\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/ir/transform.py:238\u001b[0m, in \u001b[0;36m__call__\u001b[0;34m()\u001b[0m\n\u001b[1;32m    224\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, mod):\n\u001b[1;32m    225\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Execute the pass. Note that for sequential pass, the dependency among\u001b[39;00m\n\u001b[1;32m    226\u001b[0m \u001b[38;5;124;03m    different passes will be resolved in the backend.\u001b[39;00m\n\u001b[1;32m    227\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    236\u001b[0m \u001b[38;5;124;03m        The updated module after applying this pass.\u001b[39;00m\n\u001b[1;32m    237\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 238\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _ffi_transform_api\u001b[38;5;241m.\u001b[39mRunPass(\u001b[38;5;28mself\u001b[39m, mod)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/ffi/cython/function.pxi:228\u001b[0m, in \u001b[0;36mtvm.ffi.core.Function.__call__\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:576\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    574\u001b[0m \n\u001b[1;32m    575\u001b[0m TVM_REGISTER_GLOBAL(\"transform.RunPass\")\n\u001b[0;32m--> 576\u001b[0m     .set_body_typed([](Pass pass, ffi::RValueRef<IRModule> mod) { return pass(*std::move(mod)); });\n\u001b[1;32m    577\u001b[0m \n\u001b[1;32m    578\u001b[0m TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:297\u001b[0m, in \u001b[0;36mtvm::transform::Pass::operator()(tvm::IRModule) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    295\u001b[0m \n\u001b[1;32m    296\u001b[0m IRModule Pass::operator()(IRModule mod) const {\n\u001b[0;32m--> 297\u001b[0m   return this->operator()(std::move(mod), PassContext::Current());\n\u001b[1;32m    298\u001b[0m }\n\u001b[1;32m    299\u001b[0m \n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:313\u001b[0m, in \u001b[0;36mtvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    311\u001b[0m   ret = Pass::AssertImmutableModule(mod, node, pass_ctx);\n\u001b[1;32m    312\u001b[0m } else {\n\u001b[0;32m--> 313\u001b[0m   ret = node->operator()(std::move(mod), pass_ctx);\n\u001b[1;32m    314\u001b[0m }\n\u001b[1;32m    315\u001b[0m pass_ctx.InstrumentAfterPass(ret, pass_info);\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:521\u001b[0m, in \u001b[0;36mtvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    519\u001b[0m \n\u001b[1;32m    520\u001b[0m     } else {\n\u001b[0;32m--> 521\u001b[0m       mod = pass(std::move(mod), pass_ctx);\n\u001b[1;32m    522\u001b[0m     }\n\u001b[1;32m    523\u001b[0m   }\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:313\u001b[0m, in \u001b[0;36mtvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    311\u001b[0m   ret = Pass::AssertImmutableModule(mod, node, pass_ctx);\n\u001b[1;32m    312\u001b[0m } else {\n\u001b[0;32m--> 313\u001b[0m   ret = node->operator()(std::move(mod), pass_ctx);\n\u001b[1;32m    314\u001b[0m }\n\u001b[1;32m    315\u001b[0m pass_ctx.InstrumentAfterPass(ret, pass_info);\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/ir/transform.cc:419\u001b[0m, in \u001b[0;36mtvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const\u001b[0;34m()\u001b[0m\n\u001b[1;32m    417\u001b[0m VLOG(0) << \"Executing module pass with opt level: \" << pass_info->opt_level;\n\u001b[1;32m    418\u001b[0m \n\u001b[0;32m--> 419\u001b[0m mod = pass_func(std::move(mod), pass_ctx);\n\u001b[1;32m    420\u001b[0m \n\u001b[1;32m    421\u001b[0m ICHECK(mod.defined()) << \"The return value of a module pass must be set.\";\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/tir/transforms/make_packed_api.cc:424\u001b[0m, in \u001b[0;36moperator()\u001b[0;34m()\u001b[0m\n\u001b[1;32m    422\u001b[0m }\n\u001b[1;32m    423\u001b[0m \n\u001b[0;32m--> 424\u001b[0m func = MakePackedAPI(std::move(func));\n\u001b[1;32m    425\u001b[0m \n\u001b[1;32m    426\u001b[0m if (!func.same_as(orig_func)) {\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/src/tir/transforms/make_packed_api.cc:387\u001b[0m, in \u001b[0;36mtvm::tir::MakePackedAPI(tvm::tir::PrimFunc)\u001b[0;34m()\u001b[0m\n\u001b[1;32m    385\u001b[0m \n\u001b[1;32m    386\u001b[0m   Array<Var> undefined = UndefinedVars(func_ptr->body, func_ptr->params);\n\u001b[0;32m--> 387\u001b[0m   ICHECK_EQ(undefined.size(), 0) << \"In PrimFunc \" << name_hint << \" variables \" << undefined\n\u001b[1;32m    388\u001b[0m                                  << \" are used, but are not passed in as API arguments\";\n\u001b[1;32m    389\u001b[0m \n",
      "\u001b[0;31mInternalError\u001b[0m: Check failed: undefined.size() == 0 (1 vs. 0) : In PrimFunc main variables [elem_offset] are used, but are not passed in as API arguments"
     ]
    }
   ],
   "source": [
    "check_target(\"llvm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
